import asyncio

from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain.memory import RedisChatMessageHistory
from langchain_core.runnables import RunnableWithMessageHistory

from src.common import clean_proxy
from src.module.TongyiModel import TongyiModel
from src.module.Embedding import Embedding

## 清除本地代理，否则无法连接mcp服务器
clean_proxy()

class AgentMcp:
    def __init__(self):
        print('----> 正在连接 mcp 服务器')
        self.client = asyncio.run(self.get_mcp_tools())
        self.tools = self.client.get_tools()
        print('----> mcp 服务器连接成功，检测到以下工具：', self.tools)

        self.model = TongyiModel().model
        self.embedding = Embedding().embedding
        self.prompt = ChatPromptTemplate.from_messages([
            SystemMessage(content="""你现在是一名 agent 智能体，用户需要向你进行提问，你需要结合消息的上下文历史来分析用户的问题，并使用工具来回答用户。
                如果你利仍然无法回复用户的问题，那么你就直接回复 “抱歉，我不知道问题的答案”。
            """),
            # MessagesPlaceholder(variable_name="chat_history"),
            ('human', '{input}'),
            MessagesPlaceholder(variable_name="agent_scratchpad")
        ])

    @staticmethod
    async def get_mcp_tools() -> MultiServerMCPClient:
        async with MultiServerMCPClient({
            "mcp_tools": {
                "url": 'http://127.0.0.1:8000/sse',
                "transport": 'sse'
            }
        }) as client:
            return client

    @staticmethod
    def get_message_history(session_id: str):
        return RedisChatMessageHistory(
            session_id=session_id,
            url="redis://localhost:6379/0"
        )

    def start(self):
        base_agent = create_tool_calling_agent(self.model, self.tools, self.prompt)
        agent = AgentExecutor(
            agent = base_agent,
            tools = self.tools,
            verbose = False,
            handle_parsing_errors='发生了一个错误'
        )
        # agent_with_message_history = RunnableWithMessageHistory(
        #     agent,
        #     self.get_message_history,
        #     input_messages_key='input', ## 需要固定，表示问题的字段名称，不可更换
        #     history_messages_key='chat_history' # 需要固定，此字段表示历史字段的占位符
        # )

        response = agent.invoke(
            { 'input': '深圳今天适合开车出去游玩吗?' },
        )
        print(response)